#ifndef CLSSEGMENT
#define CLSSEGMENT

#include <highgui.h>
#include <cv.h>
#include <cvaux.h>

#include "plotAPI.h"

class clsSegment
{

private:

  IplImage *segmented;         // contour edges
  IplImage *tmp;     // thresholded image
  IplImage *gray;
  IplImage *edge;
  IplImage* colorEdges;
  IplImage* dst;
  IplImage* hsv;

  // COntour fitting
  CvMemStorage* storage;
  CvPoint pt0, pt;

  CvBox2D box;
  CvPoint2D32f box_vtx[4];

  //Canny
  int minTh, maxTh;

public:

  void init(IplImage* src) //Size of image
  {
    colorEdges = cvCreateImage( cvGetSize(src), IPL_DEPTH_8U, 3 );
    hsv = cvCreateImage( cvGetSize(src), IPL_DEPTH_8U, 3 );
    edge = cvCreateImage(cvGetSize(src), IPL_DEPTH_8U, 1);
    tmp = cvCreateImage (cvGetSize(src),IPL_DEPTH_8U,3);
    gray = cvCreateImage(cvGetSize(src), IPL_DEPTH_8U, 1);
    segmented = cvCreateImage(cvGetSize(src), IPL_DEPTH_8U, 3);

    storage = cvCreateMemStorage(0);
  }

  void reset (IplImage* src)
  {
    cvClearMemStorage( storage );
    storage = cvCreateMemStorage(0);
    cvReleaseImage( &segmented );
    segmented = cvCreateImage(cvGetSize(src), IPL_DEPTH_8U, 3);
  }
  
  void release()
  {
    cvClearMemStorage( storage );
    cvReleaseImage( &colorEdges);
    cvReleaseImage( &hsv );
    cvReleaseImage( &edge );
    cvReleaseImage( &tmp );
    cvReleaseImage( &gray );
  }
  
  void getParams (double* minH, double* maxH, int index) 
  {
    IplImage* hue = cvCreateImage(cvGetSize(hsv), IPL_DEPTH_8U, 1);
    IplImage* mask = cvCreateImage(cvGetSize(hsv), IPL_DEPTH_8U, 1);

    float hranges_arr[] = {0,180};
    float* hranges = hranges_arr;
    CvHistogram *hist = 0;
    int hdims = 16;
    //int minid, maxid;
	
    hist = cvCreateHist( 1, &hdims, CV_HIST_ARRAY, &hranges, 1 );
    
    switch (index)
      {
      case 0:
	cvSplit( hsv, hue, 0, 0, 0 );
	break;
      case 1:
	cvSplit( hsv, 0, hue, 0, 0 );
	break;
      case 2:
	cvSplit( hsv, 0, 0, hue, 0 );
	break;
      default:
	cerr << "Index used in getParams is incorrect" << endl;
	break;
      }

    cvSplit( segmented, 0, 0, mask, 0 );

    cvCalcHist( &hue, hist, 0, mask );
    
    cvMinMaxLoc(hue, minH, maxH, NULL, NULL, mask);

    //*maxH = *maxH * .9;
    //*minH = *minH * 1.1;
    //cout << "mean " << mean << ", std " << sqrt(variance) << endl;
  }

 void getHists (CvHistogram* H, CvHistogram* S, CvHistogram* V) 
  {
    IplImage* hue = cvCreateImage(cvGetSize(hsv), IPL_DEPTH_8U, 1);
    IplImage* sat = cvCreateImage(cvGetSize(hsv), IPL_DEPTH_8U, 1);
    IplImage* val = cvCreateImage(cvGetSize(hsv), IPL_DEPTH_8U, 1);

    IplImage* mask = cvCreateImage(cvGetSize(hsv), IPL_DEPTH_8U, 1);

    //    float hranges_arr[] = {0,180};
    //    float* hranges = hranges_arr;
    //    int hdims = 16;
    //    int minid, maxid;
	
    cvSplit( hsv, hue, sat, val, 0 );
    cvSplit( segmented, 0, 0, mask, 0 );

    cvCalcHist( &hue, H, 0, mask );
    cvCalcHist( &sat, S, 0, mask );
    cvCalcHist( &val, V, 0, mask );
  }

  void setParams (int low, int high)
  {
    minTh = low;
    maxTh = high;
  }

  void doubleTh(IplImage* src, IplImage* dst, double minH, double maxH)
  {
    //Search and set
    for(int iY = 0; (iY < src->height); iY++) 
      {
	for(int iX = 0; (iX < src->width); iX++)
	  {
	    if ( ((uchar *)(src->imageData))[(iX+iY* src->width)] < maxH )
	      if ( ((uchar *)(src->imageData))[(iX+iY* src->width)] > minH )
		((uchar *)(dst->imageData))[(iX+iY* src->width)] = 255;
	  }
      }
  }
   

  void process(IplImage* src, IplImage* dst, CvPoint position, int width)
  {
    storage = cvCreateMemStorage(0);

    //Security check
    if (position.x < width)
      width = position.x;
    if (position.y < width)
      width = position.y;
    
    if (position.x + width > src->width)
      width = src->width - position.x;
    if (position.y + width > src->height)
      width = src->height - position.y;
    
    cvZero (segmented);

    if (width > 2)
      {
	CvRect ROI;
	
	ROI.x = position.x-width;
	ROI.y = position.y-width;
	ROI.width = 2*width;
	ROI.height = 2*width;
	
	cvSetImageROI (src, ROI);	
	cvSetImageROI (gray, ROI);
	cvSetImageROI (hsv, ROI);
	cvSetImageROI (edge, ROI);
	cvSetImageROI (colorEdges, ROI);

	cvResetImageROI (dst);

	cvCvtColor(src, gray, CV_BGR2GRAY);
	cvCvtColor(src, hsv, CV_BGR2HSV);
    
	// Run the edge detector on grayscale	 
	cvCanny(gray, edge, minTh, maxTh, 3);
	cvZero( colorEdges );
    
	cvDilate(edge,edge,NULL,5);
    
	// copy edge points
	cvCopy( hsv, colorEdges, edge );
    
	//Params
	int posX = position.x; //src->width/2;
	int posY = position.y;  //src->height/2;

	// Fitting
	CvSeq* ptseq = cvCreateSeq( CV_SEQ_KIND_GENERIC|CV_32SC2, sizeof(CvContour),
				    sizeof(CvPoint), storage );
	uchar R, G, B;
    
	R = 0;
	G = 0;
	B = 0;
    
	//Init to 0
	for(int iY = posY-width; (iY < posY+width); iY++) 
	  {
	    for(int iX = posX-width; (iX < posX+width); iX++)
	      {
		((uchar *)(dst->imageData))[(iX+iY* src->width)*3] = (uchar)(0);
		((uchar *)(dst->imageData))[(iX+iY* src->width)*3+1] = (uchar)(0);
		((uchar *)(dst->imageData))[(iX+iY* src->width)*3+2] = (uchar)(0);
	    
		((uchar *)(segmented->imageData))[(iX+iY* src->width)*3] = (uchar)(0);
		((uchar *)(segmented->imageData))[(iX+iY* src->width)*3+1] = (uchar)(0);
		((uchar *)(segmented->imageData))[(iX+iY* src->width)*3+2] = (uchar)(0);
	      }
	  }
    
	// seeding!!
	for(int iY = posY-1; (iY < posY+1); iY++) 
	  {
	    for(int iX = posX-1; (iX < posX+1); iX++)
	      {
		((uchar *)(segmented->imageData))[(iX+iY* src->width)*3] = (uchar)(255);
		((uchar *)(segmented->imageData))[(iX+iY* src->width)*3+1] = (uchar)(255);
		((uchar *)(segmented->imageData))[(iX+iY* src->width)*3+2] = (uchar)(0);
	      }
	  }   
    
	int totalPoints = -1;
    
	for (int ii=0;ii<width-1;ii++)
	  {
	    cvDilate(segmented,segmented,NULL,1);
	
	    int points = 0;
	
	    for(int iY = posY-width; (iY < posY+width); iY++) 
	      {
		for(int iX = posX+width; (iX > posX-width); iX--)
		  {		
		    if ( ((uchar *)(edge->imageData))[(iX+iY*src->width)] == 255 )
		      {
			((uchar *)(segmented->imageData))[(iX+iY* src->width)*3] = (uchar)(0);
			((uchar *)(segmented->imageData))[(iX+iY* src->width)*3+1] = (uchar)(0);
			((uchar *)(segmented->imageData))[(iX+iY* src->width)*3+2] = (uchar)(0);
		      }
		
		    R = ((uchar *)(segmented->imageData))[(iX+iY*src->width)*3];     	    
		    G = ((uchar *)(segmented->imageData))[(iX+iY*src->width)*3+1];
		
		    if ( (R==255) & (G==255) )
		      {
			((uchar *)(segmented->imageData))[(iX+iY*src->width)*3+2]=255;
		    
			points ++;
		      }
		  }
	      }
	
	
	    if ( (totalPoints == points) & (points > 0) )
	      break;
	    else
	      totalPoints = points;
	  }
    
	//defining contour
	for(int iY = posY-width; (iY < posY+width); iY++) 
	  {
	    for(int iX = posX+width; (iX > posX-width); iX--)
	      {		
		if ( ((uchar *)(segmented->imageData))[(iX+iY*src->width)*3+2] == 255 )
		  {
		    pt0.x = iX;
		    pt0.y = iY;
		
		    cvSeqPush( ptseq, &pt0 );
		  }
	      }
	  }
    
	//Clean up segmentation
	for(int iY = posY-width; (iY < posY+width); iY++) 
	  {
	    for(int iX = posX+width; (iX > posX-width); iX--)
	      {
	    
		R = ((uchar *)(segmented->imageData))[(iX+iY*src->width)*3];		    
		G = ((uchar *)(segmented->imageData))[(iX+iY*src->width)*3+1];
		B = ((uchar *)(segmented->imageData))[(iX+iY*src->width)*3+2];
	    
		if ((B == 255)) 
		  {
		    ((uchar *)(dst->imageData))[(iX+iY* src->width)*3] = (uchar)(0);
		    ((uchar *)(dst->imageData))[(iX+iY* src->width)*3+1] = (uchar)(155);
		    ((uchar *)(dst->imageData))[(iX+iY* src->width)*3+2] = (uchar)(155);
		  }
		else
		  {
		    ((uchar *)(dst->imageData))[(iX+iY* src->width)*3] = ((uchar *)(src->imageData))[(iX+iY* src->width)*3];
		    ((uchar *)(dst->imageData))[(iX+iY* src->width)*3+1] = ((uchar *)(src->imageData))[(iX+iY* src->width)*3+1];
		    ((uchar *)(dst->imageData))[(iX+iY* src->width)*3+2] =  ((uchar *)(src->imageData))[(iX+iY* src->width)*3+2];
		  }
	      }
	  }
    
	//Fit
	if (ptseq->total > 0 )
	  {
	    box = cvMinAreaRect2( ptseq, 0 );
	    cvBoxPoints( box, box_vtx );
	
	    pt0.x = cvRound(box_vtx[3].x);
	    pt0.y = cvRound(box_vtx[3].y);
	
	    for( int i = 0; i < 4; i++ )
	      {
		pt.x = cvRound(box_vtx[i].x);
		pt.y = cvRound(box_vtx[i].y);
	    
		cvLine(dst, pt0, pt, CV_RGB(0, 255, 0), 2, CV_AA, 0);
		pt0 = pt;
	      }
	  }

	cvSetImageROI (dst, ROI);
	cvResetImageROI(src);
	cvResetImageROI (gray);
	cvResetImageROI (hsv);
	cvResetImageROI (edge);
	cvResetImageROI (colorEdges);

	// Draw square
	draw_square(RED,width, dst, width, width);
      }
  }


};

#endif
